"""
高斯定理示意图
展示闭合曲面、电荷和电场线的关系
"""

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.patches import Circle
import matplotlib.patches as mpatches

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False

# 创建图形
fig = plt.figure(figsize=(16, 12))

# ===== 1. 球面对称情况（上左） =====
ax1 = fig.add_subplot(221, projection='3d')

# 球面参数
u = np.linspace(0, 2 * np.pi, 50)
v = np.linspace(0, np.pi, 50)
R = 2.0
x_sphere = R * np.outer(np.cos(u), np.sin(v))
y_sphere = R * np.outer(np.sin(u), np.sin(v))
z_sphere = R * np.outer(np.ones(np.size(u)), np.cos(v))

# 绘制球面（高斯面）
ax1.plot_surface(x_sphere, y_sphere, z_sphere, alpha=0.3, color='lightblue', edgecolor='blue', linewidth=0.5)

# 电荷（球心）
ax1.scatter([0], [0], [0], color='red', s=500, marker='o', label='Point Charge Q')
ax1.text(0, 0, 0.3, 'Q', fontsize=14, fontweight='bold', ha='center', color='red')

# 电场线（从电荷向外辐射）
theta_lines = np.linspace(0, 2*np.pi, 12)
phi_lines = np.linspace(0, np.pi, 6)
for theta in theta_lines[::2]:
    for phi in phi_lines[::2]:
        r = np.linspace(0.3, R*1.2, 20)
        x_line = r * np.sin(phi) * np.cos(theta)
        y_line = r * np.sin(phi) * np.sin(theta)
        z_line = r * np.cos(phi)
        ax1.plot(x_line, y_line, z_line, 'b-', linewidth=1.5, alpha=0.6)

# 标注
ax1.text(R*1.3, 0, 0, r'$S$', fontsize=14, fontweight='bold', color='blue')
ax1.set_xlabel('x', fontsize=11)
ax1.set_ylabel('y', fontsize=11)
ax1.set_zlabel('z', fontsize=11)
ax1.set_title('(a) Spherical Symmetry', fontsize=14, fontweight='bold')
ax1.set_xlim([-R*1.5, R*1.5])
ax1.set_ylim([-R*1.5, R*1.5])
ax1.set_zlim([-R*1.5, R*1.5])

# ===== 2. 多个电荷情况（上右） =====
ax2 = fig.add_subplot(222, projection='3d')

# 椭球面参数（高斯面）
u = np.linspace(0, 2 * np.pi, 50)
v = np.linspace(0, np.pi, 50)
a, b, c = 2.5, 1.8, 2.0
x_ellipsoid = a * np.outer(np.cos(u), np.sin(v))
y_ellipsoid = b * np.outer(np.sin(u), np.sin(v))
z_ellipsoid = c * np.outer(np.ones(np.size(u)), np.cos(v))

# 绘制椭球面
ax2.plot_surface(x_ellipsoid, y_ellipsoid, z_ellipsoid, alpha=0.3, color='lightgreen', 
                 edgecolor='green', linewidth=0.5)

# 多个电荷
charges_pos = [[-0.8, 0.5, 0.3], [0.6, -0.4, 0.2], [0.2, 0.7, -0.5]]
charges_signs = [1, -1, 1]  # 正负电荷
for pos, sign in zip(charges_pos, charges_signs):
    color = 'red' if sign > 0 else 'blue'
    marker = '+' if sign > 0 else 'o'
    ax2.scatter([pos[0]], [pos[1]], [pos[2]], color=color, s=400, marker=marker)
    label = f'Q{charges_pos.index(pos)+1}'
    ax2.text(pos[0], pos[1], pos[2]+0.3, label, fontsize=10, fontweight='bold', ha='center', color=color)

# 电场线（简化表示）
for pos, sign in zip(charges_pos, charges_signs):
    color = 'red' if sign > 0 else 'blue'
    # 从电荷发出几条电场线
    for i in range(4):
        theta = i * np.pi / 2
        phi = np.pi / 4
        r = np.linspace(0.2, 1.5, 15)
        x_line = pos[0] + r * np.sin(phi) * np.cos(theta)
        y_line = pos[1] + r * np.sin(phi) * np.sin(theta)
        z_line = pos[2] + r * np.cos(phi)
        ax2.plot(x_line, y_line, z_line, color=color, linewidth=1.2, alpha=0.5, linestyle='--')

ax2.text(a*1.2, 0, 0, r'$S$', fontsize=14, fontweight='bold', color='green')
ax2.set_xlabel('x', fontsize=11)
ax2.set_ylabel('y', fontsize=11)
ax2.set_zlabel('z', fontsize=11)
ax2.set_title('(b) Multiple Charges', fontsize=14, fontweight='bold')
ax2.set_xlim([-a*1.3, a*1.3])
ax2.set_ylim([-b*1.3, b*1.3])
ax2.set_zlim([-c*1.3, c*1.3])

# ===== 3. 2D视图：球面对称（下左） =====
ax3 = fig.add_subplot(223)

# 绘制圆形高斯面
circle = Circle((0, 0), R, fill=False, edgecolor='blue', linewidth=2.5, linestyle='-')
ax3.add_patch(circle)
ax3.plot([0], [0], 'ro', markersize=15, label='Point Charge Q')

# 电场线（径向）
for angle in np.linspace(0, 2*np.pi, 16):
    r = np.linspace(0.3, R*1.3, 30)
    x_line = r * np.cos(angle)
    y_line = r * np.sin(angle)
    ax3.plot(x_line, y_line, 'b-', linewidth=1.5, alpha=0.6)

# 标注电场方向
for angle in [0, np.pi/2, np.pi, 3*np.pi/2]:
    r = R * 0.7
    dx = 0.3 * np.cos(angle)
    dy = 0.3 * np.sin(angle)
    ax3.arrow(r*np.cos(angle), r*np.sin(angle), dx, dy, 
             head_width=0.15, head_length=0.2, fc='blue', ec='blue', linewidth=2)

# 标注
ax3.text(R*0.7, R*0.7, r'$\mathbf{E}$', fontsize=12, fontweight='bold', color='blue')
ax3.text(R*1.15, 0, r'$S$', fontsize=14, fontweight='bold', color='blue')
ax3.text(0.2, 0.2, r'$Q$', fontsize=12, fontweight='bold', color='red')
ax3.set_xlabel('x', fontsize=12, fontweight='bold')
ax3.set_ylabel('y', fontsize=12, fontweight='bold')
ax3.set_title('(c) 2D View: Spherical Symmetry', fontsize=14, fontweight='bold')
ax3.set_aspect('equal')
ax3.set_xlim([-R*1.6, R*1.6])
ax3.set_ylim([-R*1.6, R*1.6])
ax3.grid(True, alpha=0.3)
ax3.legend(loc='upper right', fontsize=10)

# ===== 4. 2D视图：多个电荷（下右） =====
ax4 = fig.add_subplot(224)

# 绘制椭圆形高斯面
ellipse = mpatches.Ellipse((0, 0), 2*a, 2*b, fill=False, edgecolor='green', linewidth=2.5, linestyle='-')
ax4.add_patch(ellipse)

# 多个电荷
for i, (pos, sign) in enumerate(zip(charges_pos, charges_signs)):
    color = 'red' if sign > 0 else 'blue'
    marker = '+' if sign > 0 else 'o'
    ax4.plot(pos[0], pos[1], marker=marker, color=color, markersize=15, markeredgewidth=2)
    label = f'Q{i+1}'
    ax4.text(pos[0], pos[1]+0.3, label, fontsize=10, fontweight='bold', ha='center', color=color)

# 电场线（简化）
for pos, sign in zip(charges_pos, charges_signs):
    color = 'red' if sign > 0 else 'blue'
    for angle in np.linspace(0, 2*np.pi, 8):
        r = np.linspace(0.2, 1.2, 20)
        x_line = pos[0] + r * np.cos(angle)
        y_line = pos[1] + r * np.sin(angle)
        ax4.plot(x_line, y_line, color=color, linewidth=1.2, alpha=0.4, linestyle='--')

# 标注
ax4.text(a*1.1, 0, r'$S$', fontsize=14, fontweight='bold', color='green')
ax4.set_xlabel('x', fontsize=12, fontweight='bold')
ax4.set_ylabel('y', fontsize=12, fontweight='bold')
ax4.set_title('(d) 2D View: Multiple Charges', fontsize=14, fontweight='bold')
ax4.set_aspect('equal')
ax4.set_xlim([-a*1.4, a*1.4])
ax4.set_ylim([-b*1.4, b*1.4])
ax4.grid(True, alpha=0.3)

# 添加图例说明
legend_elements = [
    mpatches.Patch(facecolor='red', edgecolor='red', label='Positive Charge (+)'),
    mpatches.Patch(facecolor='blue', edgecolor='blue', label='Negative Charge (-)'),
    mpatches.Patch(facecolor='none', edgecolor='blue', linewidth=2, label='Gaussian Surface S')
]
ax4.legend(handles=legend_elements, loc='upper right', fontsize=9)

plt.suptitle('Gauss\'s Law Diagram', 
             fontsize=16, fontweight='bold', y=0.98)

plt.tight_layout(rect=[0, 0, 1, 0.96])
plt.savefig('result/gauuss_law_diagram.png', dpi=300, bbox_inches='tight')
print("Gauss's Law diagram saved: gauss_law_diagram.png")
plt.close()

